"""
基于因子的债券选择策略
策略逻辑：
1. 选择"当日交易金额"因子>0.1，且"中债-模型价差"因子值为正的债券
2. 从筛选出的债券中，选取"中债-模型价差"因子值最高的10个（不足10个则全部购买）
3. 对选出的债券使用等权重配置（目标仓位为1/N）
"""
from datetime import datetime
from typing import Dict, List, Optional
import pandas as pd
import numpy as np
from loguru import logger

from strategy.base import Strategy
from engine.object import BarData, FactorData
from engine.constant import AssetType

class BondFactorStrategy(Strategy):
    """基于因子的债券选择策略"""
    
    def __init__(self, name: str = "BondFactorStrategy", 
                 leverage: float = 1.0,
                 rebalance_interval: int = 1,  # 每隔5个交易日重新平衡一次
                 portfolio_flag: bool = True,
                 score_floor=0.025,
                 score_floor_buffer=0.0,  # 缓冲池的因子值阈值，默认为0
                 score_cap=0.05,
                 category:Optional[str]=None,
                 ):
        """
        初始化策略
        
        Args:
            name: 策略名称
            initial_capital: 初始资金
            risk_limit: 风险限额，以总资产比例表示
            leverage: 杠杆水平
            rebalance_interval: 再平衡间隔（交易日）
            category: 债券类别,"国债"或者"政策性金融债"或者None
        """
        super().__init__(name,  leverage,portfolio_flag)
        self.rebalance_interval = rebalance_interval
        self.day_counter = 0
        self.selected_bonds = []  # 当前选中的债券列表
        self.score_floor_buffer = score_floor_buffer  # 缓冲池的因子值阈值
        self.score_floor = score_floor
        self.score_cap = score_cap
        self.category = category
        logger.info(f"基于因子的债券选择策略 {name} 初始化完成")
        
        # 注意：不要在初始化时就加载数据，因为此时engine还未设置
    
    def initialize(self) -> None:
        """
        策略初始化
        加载所有可用标的和因子数据
        """
        pass
    
    def on_dataset(self, data: Dict[str, BarData]) -> None:
        """
        数据更新事件，处理最新市场数据，实现选股逻辑
        
        Args:
            data: 最新市场数据
        """
        # 更新当前日期

        self.current_date = self.engine.current_time

        
        # 每隔指定的交易日数重新选择债券
        self.day_counter += 1
        if self.day_counter >= self.rebalance_interval:
            self.day_counter = 0
            self._select_bonds()
    
    def _select_bonds(self):
        """
        基于因子值选择债券，包含缓冲池机制
        选择逻辑：
        1. 选择"当日交易金额"因子>0.1，且"中债-模型价差"因子值大于score_floor的债券作为主要候选
        2. 对于已在持仓中的债券，如果其因子值>score_floor_buffer但<score_floor，则保留在缓冲池中
        3. 从筛选出的债券中，选取"中债-模型价差"因子值最高的债券
        4. 对选出的债券使用等权重配置
        """
        # target category
        if self.category is not None:
            effective_symbols = [i for i  in self.engine.assets if \
                                        self.engine.assets[i].secucategory == self.category]
        else:
            effective_symbols = list(self.engine.assets.keys())
        # 获取当前持仓信息
        current_positions = self.portfolio.get_position_details()
        current_holding_symbols = [pos.symbol for pos in current_positions if pos.quantity != 0]
        
        # 获取当天所有标的的所有因子值
        all_factors = self.engine.get_today_factor_value()
        if not all_factors:
            logger.warning(f"当前日期 {self.current_date} 没有因子数据")
            return
        
        # 获取特定因子值
        volume_factors = self.engine.get_today_factor_value("当日交易金额")
        price_diff_factors = self.engine.get_today_factor_value("中债-模型价差")
        
        # 主要候选池：满足主要筛选条件的债券
        main_pool = []
        main_pool_values = []
        
        # 缓冲池：已在持仓中且因子值在缓冲区间的债券
        buffer_pool = []
        buffer_pool_values = []
        
        # 筛选符合条件的债券
        for symbol in all_factors.keys():

            # 检查标的是否有相关因子数据，如果没有直接跳过
            if symbol not in volume_factors or \
                symbol not in price_diff_factors or \
                symbol not in effective_symbols:
                continue
            
            # 获取因子值
            volume = volume_factors[symbol]
            price_diff = price_diff_factors[symbol]
    
            # 主要筛选条件：交易金额>0.1 且 模型价差>score_floor 且 <score_cap
            if volume > 0.1 and price_diff > self.score_floor and price_diff < self.score_cap:
                main_pool.append(symbol)
                main_pool_values.append(price_diff)
            # 缓冲池条件：已在持仓中 且 模型价差>score_floor_buffer 但 <score_floor
            elif symbol in current_holding_symbols and volume > 0.1 and \
                 price_diff > self.score_floor_buffer and price_diff < self.score_floor:
                buffer_pool.append(symbol)
                buffer_pool_values.append(price_diff)
        
        # 如果主要候选池和缓冲池都为空，保持当前持仓
        if not main_pool and not buffer_pool:
            logger.info("没有符合条件的债券，保持当前持仓")
            return
        
        # 合并主要候选池和缓冲池
        combined_pool = main_pool + buffer_pool
        combined_values = main_pool_values + buffer_pool_values

        
        # 对合并后的候选池按照价差大小排序
        sorted_indices = np.argsort(combined_values)[::-1]  # 降序排序
        
        # 取前15个（如不足15个则全部选择）
        top_n = min(15, len(combined_pool))
        final_selected = [combined_pool[i] for i in sorted_indices[:top_n]]
        
        # 记录选择的债券
        self.selected_bonds = final_selected
        
        # 记录缓冲池使用情况
        buffer_count = sum(1 for bond in final_selected if bond in buffer_pool)
        logger.info(f"选择了 {len(final_selected)} 只债券，其中 {buffer_count} 只来自缓冲池")
        if buffer_count > 0:
            buffer_bonds = [bond for bond in final_selected if bond in buffer_pool]
            logger.info(f"缓冲池债券: {buffer_bonds}")
        
        # 计算目标权重（等权重配置）
        target_weights = {}
        weight_per_bond = 1.0 / len(final_selected)
        for symbol in final_selected:
            target_weights[symbol] = weight_per_bond

        # 调整目标仓位
        self.adjust_target_position(target_weights)
        logger.info(f"目标权重设置为: {target_weights}")
